package org.rbio.paper.methylglm.annotate;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.broadinstitute.sting.commandline.Argument;
import org.broadinstitute.sting.commandline.RodBinding;
import org.broadinstitute.sting.commandline.Tags;
import org.broadinstitute.sting.gatk.GenomeAnalysisEngine;
import org.broadinstitute.sting.gatk.contexts.AlignmentContext;
import org.broadinstitute.sting.gatk.contexts.ReferenceContext;
import org.broadinstitute.sting.gatk.datasources.rmd.ReferenceOrderedDataSource;
import org.broadinstitute.sting.gatk.refdata.RefMetaDataTracker;
import org.broadinstitute.sting.gatk.refdata.tracks.RMDTrackBuilder;
import org.broadinstitute.sting.gatk.refdata.utils.RMDTriplet;
import org.broadinstitute.sting.gatk.walkers.RodWalker;
import org.broadinstitute.sting.utils.codecs.refseq.RefSeqFeature;
import org.broadinstitute.sting.utils.codecs.table.TableFeature;

// CpG annotation class [GATK walker]
// annotates: TSS, gene, exon/intorn, CpG islands and CpG shores
// 1 based positions
public class AnnotateCpGWalker extends RodWalker<Integer, Integer> {
	private final static String NA = "NA";

	// program arguments
	@Argument(shortName = "ref_fa", doc = "reference fasta", required = true)
	private String ref_fa;

	@Argument(shortName = "ucsc_gene", doc = "UCSC gene table", required = true)
	private String ucsc_gene;

	@Argument(shortName = "cpg_isl", doc = "UCSC CpG island table", required = true)
	private String cpg_isl;

	@Argument(shortName = "meth_file", doc = "CpG deviance file", required = true)
	private String meth_file;

	@Argument(shortName = "dmr_file", doc = "DMR file", required = true)
	private String dmr_file;

	@Argument(shortName = "tss_up", doc = "bp upstream of TSS", required = true)
	private int tss_up;

	@Argument(shortName = "tss_down", doc = "bp downstream of TSS", required = true)
	private int tss_down;

	@Argument(shortName = "out_prefix", doc = "output prefix", required = true)
	private String out_prefix;

	// CpG rod file
	RodBinding<TableFeature> cpg_rod;

	// CpG island file
	RodBinding<TableFeature> cpg_isl_rod;

	// TSS file
	RodBinding<TableFeature> tss_rod;

	// UCSC gene file
	RodBinding<RefSeqFeature> gene_rod;

	// used to print the header
	boolean is_first_cpg = true;
	BufferedWriter obuf_cpg;
	String tmp_prefix;

	// initialize
	@Override
	public void initialize() {
		try {
			// CpG annotation file
			obuf_cpg = new BufferedWriter(new FileWriter(out_prefix
					+ ".cpg.ann.txt"));

			// initialize the annotation files
			tmp_prefix = PrepareAnnotationTables.prepare(ref_fa, ucsc_gene,
					cpg_isl, meth_file, tss_up, tss_down);

			// add the RODs
			GenomeAnalysisEngine engine = this.getToolkit();
			RMDTrackBuilder builder = new RMDTrackBuilder(
					engine.getMasterSequenceDictionary(),
					engine.getGenomeLocParser(), null);

			// CpG island
			RMDTriplet fdesc = new RMDTriplet("cpg_isl", "table", tmp_prefix
					+ ".isl.gatk.txt", RMDTriplet.RMDStorageType.FILE,
					new myTags("table"));
			engine.getRodDataSources().add(
					new ReferenceOrderedDataSource(fdesc, builder, engine
							.getMasterSequenceDictionary(), engine
							.getGenomeLocParser(), false));

			// CpGs
			fdesc = new RMDTriplet("cpg", "table",
					tmp_prefix + ".cpg.gatk.txt",
					RMDTriplet.RMDStorageType.FILE, new myTags("table"));
			engine.getRodDataSources().add(
					new ReferenceOrderedDataSource(fdesc, builder, engine
							.getMasterSequenceDictionary(), engine
							.getGenomeLocParser(), false));

			// gene
			fdesc = new RMDTriplet("gene", "refseq", tmp_prefix
					+ ".ucsc.gatk.txt", RMDTriplet.RMDStorageType.FILE,
					new myTags("refseq"));
			engine.getRodDataSources().add(
					new ReferenceOrderedDataSource(fdesc, builder, engine
							.getMasterSequenceDictionary(), engine
							.getGenomeLocParser(), false));

			// tss
			fdesc = new RMDTriplet("tss", "table",
					tmp_prefix + ".tss.gatk.txt",
					RMDTriplet.RMDStorageType.FILE, new myTags("table"));
			engine.getRodDataSources().add(
					new ReferenceOrderedDataSource(fdesc, builder, engine
							.getMasterSequenceDictionary(), engine
							.getGenomeLocParser(), false));

			// rod objects to reference in the tracker
			cpg_isl_rod = new RodBinding<TableFeature>(TableFeature.class,
					"cpg_isl", tmp_prefix + ".isl.gatk.txt", "table",
					new myTags("table"));
			cpg_rod = new RodBinding<TableFeature>(TableFeature.class, "cpg",
					tmp_prefix + ".cpg.gatk.txt", "table", new myTags("table"));
			gene_rod = new RodBinding<RefSeqFeature>(RefSeqFeature.class,
					"gene", tmp_prefix + ".ucsc.gatk.txt", "refseq",
					new myTags("refseq"));
			tss_rod = new RodBinding<TableFeature>(TableFeature.class, "tss",
					tmp_prefix + ".tss.gatk.txt", "table", new myTags("table"));
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// header
	public void writeHeader(TableFeature cpg) {
		// header of the CpG file
		StringBuilder col_str = new StringBuilder();
		for (int i = 1; i < cpg.getHeader().size(); i++) {
			col_str.append(cpg.getHeader().get(i) + "\t");
		}
		try {
			obuf_cpg.write(String.format("%stss\tgene\texon\tisl\n", col_str));
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// mapping routine
	@Override
	public Integer map(RefMetaDataTracker tracker, ReferenceContext ref,
			AlignmentContext context) {
		if (tracker == null) {
			return null;
		}

		// features
		TableFeature cpg = tracker.getFirstValue(cpg_rod);
		TableFeature isl = tracker.getFirstValue(cpg_isl_rod);
		List<TableFeature> tss_lst = tracker.getValues(tss_rod);
		List<RefSeqFeature> tx_lst = tracker.getValues(gene_rod);

		// process the CpGs
		if (cpg != null) {
			// default NA values
			String tss_str = NA, gene = NA, gene_body_type = NA, isl_str = NA;

			// check for isl [there can be only 1 island for a CpG]
			if (isl != null) {
				isl_str = "YES";
			}

			// check for tss [tss can overlap]
			if (tss_lst.size() > 0) {
				tss_str = tss_lst.get(0).get("gene");
				for (int i = 1; i < tss_lst.size(); i++) {
					tss_str += "," + tss_lst.get(i).get("gene");
				}
			}

			// check for gene [genes can overlap]
			if (tx_lst.size() > 0) {
				// make gene-tx multimap
				// allows for several genes with each gene having several
				// transcripts
				Map<String, List<RefSeqFeature>> gene_tab = new TreeMap<String, List<RefSeqFeature>>();
				for (RefSeqFeature tx : tx_lst) {
					List<RefSeqFeature> gene_tx_lst = gene_tab.get(tx
							.getGeneName());
					if (gene_tx_lst == null) {
						gene_tx_lst = new ArrayList<RefSeqFeature>();
						gene_tab.put(tx.getGeneName(), gene_tx_lst);
					}
					gene_tx_lst.add(tx);
				}

				// check if the CpG is in the exon or intron of a gene.
				// for each gene we check if the CpG is in any exon of the
				// transcript. If such an exon exist then report that CpG
				// as in the EXON for this gene o/w report INTRON for this
				// gene.
				// each gene will have only 1 gene_body_type set : EXON or
				// INTRON
				List<String> gene_lst = new ArrayList<String>();
				List<String> gene_body_type_lst = new ArrayList<String>();
				for (String gene_name : gene_tab.keySet()) {
					gene_lst.add(gene_name);
					String cur_gene_body_type = "INTRON";
					for (RefSeqFeature tx : gene_tab.get(gene_name)) {
						if (tx.overlapsExonP(ref.getLocus()) == true) {
							cur_gene_body_type = "EXON";
							break;
						}
					}
					gene_body_type_lst.add(cur_gene_body_type);
				}

				// set the values
				gene = gene_lst.get(0);
				gene_body_type = gene_body_type_lst.get(0);
				for (int i = 1; i < gene_lst.size(); i++) {
					gene += "," + gene_lst.get(i);
					gene_body_type += "," + gene_body_type_lst.get(i);
				}
			}

			// write the header
			if (is_first_cpg == true) {
				writeHeader(cpg);
				is_first_cpg = false;
			}

			// write the annotated CpGs
			StringBuilder col_str = new StringBuilder();
			for (int i = 1; i < cpg.getAllValues().size(); i++) {
				col_str.append(cpg.getAllValues().get(i) + "\t");
			}
			try {
				obuf_cpg.write(String.format("%s%s\t%s\t%s\t%s\n", col_str,
						tss_str, gene, gene_body_type, isl_str));
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
		return null;
	}

	@Override
	// not used
	public Integer reduceInit() {
		return null;
	}

	@Override
	// not used
	public Integer reduce(Integer value, Integer sum) {
		return null;
	}

	// annoate the DMRs
	public void onTraversalDone(Integer result) {
		try {
			obuf_cpg.close();
			DMRAnnotate.annotate(out_prefix + ".cpg.ann.txt", dmr_file,
					out_prefix + ".dmr.ann.txt");

			// do the clean up
			(new File(tmp_prefix + ".isl.gatk.txt")).delete();
			(new File(tmp_prefix + ".cpg.gatk.txt")).delete();
			(new File(tmp_prefix + ".ucsc.gatk.txt")).delete();
			(new File(tmp_prefix + ".tss.gatk.txt")).delete();

		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// my tags
	public static class myTags extends Tags {
		public myTags(String type) {
			super();
			addPositionalTag(type);
		}
	}
}
